import math
import pandas as pd
import backtrader as bt
from pyLibs import BacktraderUtits as btu


class BollingerVolatilityBreakout(btu.BaseStrategy):
    params = (
        ('printlog', True),
        ('smoothing_period', 5),
        ('stack_len', 3),
    )

    def __init__(self):
        self.datasLength = len(self.getdatanames())
        self.watchSign = self.orderSign = False
        self.buyPrice = self.sellPrice = self.lossPrice = self.initial_loss = None
        self.columns = ['stock_code', 'datetime', 'open', 'close', 'high', 'low', 'ema_m', 'ema_l', 'bollinger_mid',
                        'mid_ema_position', 'platform_high', 'platform_low', 'atr', 'initial_loss', 'moving_lost',
                        'break_high_enter', 'break_low_enter', 'is_trigger_initial', 'is_trigger_moving', 'status']
        self.indicators = list()
        self.condition = list()
        for i in range(0, self.datasLength):
            if self.data._timeframe == bt.TimeFrame.Minutes:
                if self.data._compression in [5, 15, 30, 60]:
                    ema_h34 = int(60 * 34 / self.data._compression)
                    ema_d34 = int(60 * 34 * 4 / self.data._compression)
                    bollinger_h = int(60 * 20 / self.data._compression)
                    bollinger_d = int(60 * 20 * 4 / self.data._compression)
                    atr_m = int(15 * 14 / self.data._compression)
                    atr_l = int(60 * 14 / self.data._compression)
                    ema_m = bt.indicators.ExponentialMovingAverage(self.datas[i], period=ema_h34)             # 小时线 EMA(34)
                    ema_l = bt.indicators.ExponentialMovingAverage(self.datas[i], period=ema_d34)             # 日线 EMA(34)
                    boll_m = bt.indicators.BollingerBands(self.datas[i], period=bollinger_h, devfactor=2.0)   # 小时线 bollinger(20)
                    boll_l = bt.indicators.BollingerBands(self.datas[i], period=bollinger_d, devfactor=2.0)   # 日线 bollinger(20)
                    atr_m = bt.indicators.AverageTrueRange(self.datas[i], period=atr_m)                       # 15分钟线 ATR(14)
                    atr_l = bt.indicators.AverageTrueRange(self.datas[i], period=atr_l)                       # 小时线 ATR(14)
                    pass
                else:
                    pass
            elif self.data._timeframe == bt.TimeFrame.Days:
                ema_m = bt.indicators.ExponentialMovingAverage(self.datas[i], period=34)            # 日线 EMA(34)
                ema_l = bt.indicators.ExponentialMovingAverage(self.datas[i], period=170)           # 周线 EMA(34)
                boll_m = bt.indicators.BollingerBands(self.datas[i], period=20, devfactor=2.0)      # 日线 bollinger(20)
                boll_l = bt.indicators.BollingerBands(self.datas[i], period=100, devfactor=2.0)     # 周线 bollinger(20)
                atr_m = bt.indicators.AverageTrueRange(self.datas[i], period=3)                     # 小时线 ATR(15)
                atr_l = bt.indicators.AverageTrueRange(self.datas[i], period=14)                    # 日线 ATR(14)
                pass
            else:
                pass
            close = self.datas[i].close
            volume = self.datas[i].volume
            vol_ma_l = bt.indicators.SmoothedMovingAverage(volume, period=89)
            self.indicators.append({
                'vol_ma_s': bt.indicators.SmoothedMovingAverage(self.datas[i].volume, period=13),
                'vol_ma_l': vol_ma_l,
                'ema_m': ema_m,
                'ema_l': ema_l,
                'boll_m_mid': boll_m.mid,
                'boll_m_top': boll_m.top,
                'boll_m_bot': boll_m.bot,
                'boll_m_bb': (close - boll_m.bot) / (boll_m.top - boll_m.bot),
                'boll_m_width': (boll_m.top - boll_m.bot) / boll_m.mid,
                'boll_l_mid': boll_l.mid,
                'boll_l_top': boll_l.top,
                'boll_l_bot': boll_l.bot,
                'boll_l_bb': (close - boll_l.bot) / (boll_l.top - boll_l.bot),
                'boll_l_width': (boll_m.top - boll_l.bot) / boll_l.mid,
                'atr_m': atr_m,
                'atr_l': atr_l,
            })
            self.condition.append({
                'is_vol_release': volume > vol_ma_l,                                        # 是否放量
                'is_vol_multiplier': bt.And(volume > vol_ma_l * 2, volume > volume[-1]),    # 是否倍量
                'is_mid_m_increase': boll_m.mid > bt.talib.SMA(boll_m.mid, timeperiod=5),   # 中期布林中轨是都上升
                'is_mid_l_increase': boll_l.mid > bt.talib.SMA(boll_l.mid, timeperiod=5),   # 长期布林中轨是否上升
                'is_ema_m_increase': ema_m > bt.talib.SMA(ema_m, timeperiod=5),             # 中期EMA均线是否上升
                'is_ema_l_increase': ema_l > bt.talib.SMA(ema_l, timeperiod=5),             # 长期EMA均线是否上升
                'close_over_mid_m': bt.indicators.CrossOver(close, boll_m.mid),             # 收盘是否金叉中期布林中轨
                'close_over_mid_l': bt.indicators.CrossOver(close, boll_l.mid),             # 收盘是否金叉长期布林中轨
                'mid_m_over_ema_m': bt.indicators.CrossOver(boll_m.mid, ema_m),             # 中期布林中轨是否金叉中期EMA
                'mid_m_over_ema_l': bt.indicators.CrossOver(boll_m.mid, ema_l),             # 中期布林中轨是否金叉长期EMA
                'mid_l_over_ema_l': bt.indicators.CrossOver(boll_l.mid, ema_l),             # 长期布林中轨是否金叉长期EMA
                'ma_short_term_long': boll_m.mid > ema_m,                                   # 均线短期多头
                'ma_long_term_long': bt.And(boll_m.mid > ema_m,  ema_m > ema_l),            # 均线长期多头
            })
            self.longBuySign = bt.And(self.condition[i]['is_vol_multiplier'] > 0,
                             self.condition[i]['is_mid_m_increase'] > 0,
                             self.condition[i]['close_over_mid_m'] > 0)

    def next(self):
        for i in range(0, self.datasLength):
            stock = self.datas[i]
            indicators = self.indicators[i]
            condition = self.condition[i]
            low_min = min([round(stock.low[x], 3) for x in range(-8, 0)])
            atr_max = max([round((stock.low[x] - indicators['atr_m'][x] * 1.5), 3) for x in range(-8, 0)])
            if self.longBuySign and self.orderSign is False:
                symbol_multi = self.broker.getcommissioninfo(stock).p.mult
                total_value = self.broker.getvalue()
                lots = math.floor(total_value / (symbol_multi * stock.close[0]) / 100) * 100
                self.lossPrice = self.initial_loss = min(low_min, atr_max)
                self.buy_bracket(stock, size=lots, stopprice=self.lossPrice)
                self.log('买入 信号, 交易单价: %.2f, 交易数量: %i, 当前仓位: %i, 初始止损价：%.2f' %
                         (stock.close[0], lots, self.getposition(stock).size, self.lossPrice))
            if self.orderSign:
                if self.lossPrice < atr_max and stock.close[0] > self.buyPrice * 1.01:
                    # 向上调整移动止损位
                    self.lossPrice = atr_max
                    self.log('当前收盘价：%.2f，买入价格：%2f，进行【调整移动止损位】上调至: %.2f' % (
                        stock.close[0], self.buyPrice, atr_max))
                if stock.close[0] < self.lossPrice:
                    if self.close(stock):
                        if self.initial_loss  == self.lossPrice:
                            self.log('初始止损 信号 (平仓), 交易单价: %.2f, 交易数量: %i' % (
                                stock.close[0], self.getposition(stock).size))
                        else:
                            self.log('移动止盈 信号 (平仓), 交易单价: %.2f, 交易数量: %i' % (
                                stock.close[0], self.getposition(stock).size))